package com.jhc.yolo;

import ai.onnxruntime.OrtException;
import com.google.gson.Gson;
import com.jhc.yolo.exception.FileNotValidException;
import com.jhc.yolo.pojo.Detection;
import com.jhc.yolo.pojo.ModelConfig;
import com.jhc.yolo.pojo.Yolo;
import org.opencv.core.Mat;
import org.opencv.core.MatOfByte;
import org.opencv.imgcodecs.Imgcodecs;

import java.io.File;
import java.io.IOException;
import java.util.List;

/**
 * 识别入口
 */
public class Runner {
    private final Yolo inferenceSession;
    private static final Gson gson = new Gson();

    public Runner(ModelConfig config) throws IOException, OrtException {
        this.inferenceSession = new ModelFactory().getModel(config);
    }
    public String detection(String filePath) throws OrtException {
        fileValid(filePath);
        Mat img = Imgcodecs.imread(filePath, Imgcodecs.IMREAD_COLOR);
        List<Detection> detectionList = inferenceSession.run(img);
        return gson.toJson(detectionList);

    }

    /**
     *
     * @param filePath 图片路径
     * @param savePath 预测结果保存位置,需精确到文件名如/xxx/xxx/xxx.jpg
     * @return
     * @throws OrtException
     */
    public String detection(String filePath,String savePath) throws OrtException {
        fileValid(filePath);
        Mat img = Imgcodecs.imread(filePath, Imgcodecs.IMREAD_COLOR);
        List<Detection> detectionList = inferenceSession.run(img);
        ImageUtil.drawPredictions(img, detectionList);
        Imgcodecs.imwrite(savePath, img);
        return gson.toJson(detectionList);
    }

    /**
     * @param bytes MultipartFile.getBytes后传入即可
     * @return
     */
    public String detection(byte[] bytes) throws OrtException {
        Mat img = Imgcodecs.imdecode(new MatOfByte(bytes), Imgcodecs.IMREAD_COLOR);
        List<Detection> detectionList = inferenceSession.run(img);
        return gson.toJson(detectionList);
    }

    /**
     *
     * @param bytes MultipartFile.getBytes后传入即可
     * @param savePath 预测结果保存位置,需精确到文件名如/xxx/xxx/xxx.jpg,目录不存在则不保存
     * @return
     * @throws OrtException
     */
    public String detection(byte[] bytes,String savePath) throws OrtException {
        Mat img = Imgcodecs.imdecode(new MatOfByte(bytes), Imgcodecs.IMREAD_COLOR);
        List<Detection> detectionList = inferenceSession.run(img);
        ImageUtil.drawPredictions(img, detectionList);
        Imgcodecs.imwrite(savePath, img);
        return gson.toJson(detectionList);
    }

    // 文件合法性校验
    public void fileValid(String path){
        File file = new File(path);
        if (!file.exists()) {
            throw new FileNotValidException("文件不存在: " + path);
        }
        if (file.isDirectory()) {
            throw new FileNotValidException(path + " 不是文件");
        }
    }
}
